package net.lindeer.lang

import java.util.Stack

class Resolver(private val interpreter: Interpreter) : Expr.Visitor<Unit>,
        Stmt.Visitor<Unit> {
    private enum class FunctionType {
        NONE,
        FUNCTION,
        INITIALIZER,
        METHOD,
    }
    private enum class ClassType {
        NONE,
        CLASS,
        SUBCLASS,
    }
    private val scopes = Stack<MutableMap<String, Boolean>>()
    private var currentFunction = FunctionType.NONE
    private var currentClass = ClassType.NONE

    override fun visitBinaryExpr(expr: Expr.Binary) {
        resolveExpression(expr.left)
        resolveExpression(expr.right)
    }

    override fun visitGroupingExpr(expr: Expr.Grouping) {
        resolveExpression(expr.expression)
    }

    override fun visitLiteralExpr(expr: Expr.Literal) {
    }

    override fun visitUnaryExpr(expr: Expr.Unary) {
        resolveExpression(expr.right)
    }

    override fun visitVariableExpr(expr: Expr.Variable) {
        val token = expr.name
        if (scopes.isNotEmpty() && scopes.peek()[token.lexeme] != true) {
            Rox.error(token, "Can't read local variable in its own initializer!")
        }
        resolveLocal(expr, expr.name)
    }

    override fun visitAssignExpr(expr: Expr.Assign) {
        resolveExpression(expr.value)
        resolveLocal(expr, expr.name)
    }

    override fun visitLogicalExpr(expr: Expr.Logical) {
        resolveExpression(expr.left)
        resolveExpression(expr.right)
    }

    override fun visitCallExpr(expr: Expr.Call) {
        resolveExpression(expr.callee)
        expr.arguments.forEach {
            resolveExpression(it)
        }
    }

    override fun visitExpressionStmt(stmt: Stmt.Expression) {
        resolveExpression(stmt.expression)
    }

    override fun visitPrintStmt(stmt: Stmt.Print) {
        resolveExpression(stmt.expression)
    }

    override fun visitVarStmt(stmt: Stmt.Var) {
        declare(stmt.name)
        stmt.initializer?.let { resolveExpression(it) }
        define(stmt.name)
    }

    override fun visitBlockStmt(stmt: Stmt.Block) {
        beginScope()
        resolve(stmt.statements)
        endScope()
    }

    override fun visitIfStmt(stmt: Stmt.If) {
        resolveExpression(stmt.condition)
        resolveStatement(stmt.thenStmt)
        stmt.elseStmt?.let { resolveStatement(it) }
    }

    override fun visitWhileStmt(stmt: Stmt.While) {
        resolveExpression(stmt.condition)
        resolveStatement(stmt.body)
    }

    override fun visitFuncStmt(stmt: Stmt.Func) {
        resolveFunction(stmt, FunctionType.FUNCTION)
    }

    override fun visitReturnStmt(stmt: Stmt.Return) {
        if (currentFunction == FunctionType.NONE) {
            Rox.error(stmt.keyword, "Can't return from top-level code!")
        }
        stmt.value?.let {
            if (currentFunction == FunctionType.INITIALIZER) {
                Rox.error(stmt.keyword, "Can't return a value from initializer!")
            }
            resolveExpression(it)
        }
    }

    override fun visitClazzStmt(stmt: Stmt.Clazz) {
        val classType = currentClass
        currentClass = ClassType.CLASS
        declare(stmt.name)

        stmt.superclass?.let {
            if (stmt.name.lexeme == it.name.lexeme) {
                Rox.error(it.name, "A class can't inherit from itself!")
            }
            currentClass = ClassType.SUBCLASS
            resolveExpression(it)
            beginScope()
            scopes.peek()["super"] = true
        }

        beginScope()
        scopes.peek()["this"] = true
        stmt.methods.forEach {
            val funcType = if (it.name.lexeme == "init") FunctionType.INITIALIZER else FunctionType.METHOD
            resolveFunction(it, funcType)
        }
        define(stmt.name)
        endScope()
        stmt.superclass?.let {
            endScope()
        }
        currentClass = classType
    }

    override fun visitGetterExpr(expr: Expr.Getter) {
        resolveExpression(expr.who)
    }

    override fun visitSetterExpr(expr: Expr.Setter) {
        resolveExpression(expr.right)
        resolveExpression(expr.who)
    }

    override fun visitThisExpr(expr: Expr.This) {
        if (currentClass == ClassType.NONE) {
            Rox.error(expr.keyword, "Can't ust 'this' outside of a class!")
            return
        }
        resolveLocal(expr, expr.keyword)
    }

    override fun visitSuperExpr(expr: Expr.Super) {
        when {
            currentClass == ClassType.NONE -> {
                Rox.error(expr.keyword, "Can't use 'super' outside of a class!")
            }
            currentClass != ClassType.SUBCLASS -> {
                Rox.error(expr.keyword, "Can't use 'super' in a class without super class!")
            }
        }
        resolveLocal(expr, expr.keyword)
    }

    fun resolve(statements: List<Stmt>) {
        statements.forEach { resolveStatement(it) }
    }

    private fun resolveStatement(stmt: Stmt) {
        stmt.accept(this)
    }

    private fun resolveExpression(expr: Expr) {
        expr.accept(this)
    }

    private fun beginScope() {
        scopes.push(mutableMapOf())
    }

    private fun endScope() {
        scopes.pop()
    }

    private fun declare(name: Token) {
        if (scopes.isEmpty()) {
            return
        }
        val scope = scopes.peek()
        if (scope.containsKey(name.lexeme)) {
            Rox.error(name, "Already variable with this name in this scope!")
        }
        scope[name.lexeme] = false
    }

    private fun define(name: Token) {
        if (scopes.isEmpty()) {
            return
        }
        scopes.peek()[name.lexeme] = true
    }

    private fun resolveLocal(expr: Expr, name: Token) {
        scopes.asReversed().forEachIndexed { index, scope ->
            if (scope.containsKey(name.lexeme)) {
                interpreter.resolve(expr, index)
                return
            }
        }
    }

    private fun resolveFunction(func: Stmt.Func, type: FunctionType) {
        val enclosingFunc = currentFunction
        currentFunction = type
        beginScope()
        declare(func.name)
        define(func.name)
        func.params.forEach { param ->
            declare(param)
            define(param)
        }
        resolve(func.body)
        endScope()
        currentFunction = enclosingFunc
    }
}
